#include "../Python/PythonS2E.h"

#include <llvm/Support/FileSystem.h>
#include <llvm/Support/Path.h>
#include <s2e/ConfigFile.h>
#include <s2e/S2E.h>
#include <s2e/S2EExecutor.h>
#include <Python.h>
#include <dlfcn.h>
#include <string>
#include <stdlib.h>
#include "bindings.h"

namespace s2e {
namespace plugins {
namespace python {

S2E_DEFINE_PLUGIN(PythonS2E, "Execute Python Script", "PythonS2E",
		"ModuleExecutionDetector");

void PythonS2E::initialize() {
	ConfigFile *cfg = s2e()->getConfig();
	bool ok = false;

	m_python_so_file = cfg->getString(getConfigKey() + ".python_so_file", "",
			&ok);
	if (!ok) {
		getWarningsStream() << "Please specify PYTHON_LIBRARIES\n";
		exit(-1);
	}
	if (!llvm::sys::fs::exists(m_python_so_file)) {
		getWarningsStream() << m_python_so_file
				<< "python_so_file does not exist\n";
		exit(-1);
	}
	m_script_file = cfg->getString(getConfigKey() + ".script_file", "", &ok);
	if (!llvm::sys::fs::exists(m_script_file)) {
		getWarningsStream() << m_script_file << "script_file does not exist\n";
		exit(-1);
	}

	m_detector = s2e()->getPlugin<ModuleExecutionDetector>();
	m_monitor = static_cast<OSMonitor*>(s2e()->getPlugin("OSMonitor"));
	m_process = s2e()->getPlugin<ProcessExecutionDetector>();

	// for normal instruction instrument
	m_detector->onModuleTranslateBlockStart.connect(
			sigc::mem_fun(*this, &PythonS2E::onModuleTranslateBlockStart));

	/* This is for syscall entry monitoring */
	s2e()->getCorePlugin()->onTranslateSoftInterruptStart.connect(
			sigc::mem_fun(*this, &PythonS2E::onTranslateSoftInterruptStart));

	s2e()->getCorePlugin()->onTranslateSpecialInstructionEnd.connect(
			sigc::mem_fun(*this, &PythonS2E::onTranslateSpecialInstructionEnd));
	//for syscall exit
	s2e()->getCorePlugin()->onTranslateInstructionStart.connect(
			sigc::mem_fun(*this, &PythonS2E::onEveryTranslateInstructionStart));

	/* Init the python module */
	PyImport_AppendInittab("plugins", s2e::plugins::initplugins);

	Py_Initialize();

	if (llvm::sys::fs::exists(m_script_file)) {
		execScript(m_script_file.c_str());
	}
	return;
}

void PythonS2E::onTranslateSpecialInstructionEnd(ExecutionSignal *signal,
		S2EExecutionState *state, TranslationBlock *tb, uint64_t pc,
		enum special_instruction_t type) {
	if (type != SYSCALL && type != SYSENTER) {
		return;
	}

	uint64_t pid = m_monitor->getPid(state);
	if (!m_process->isTracked(state, pid)) {
		return;
	}

	signal->connect(sigc::mem_fun(*this, &PythonS2E::onSyscallInst));
}

void PythonS2E::onTranslateSoftInterruptStart(ExecutionSignal *signal,
		S2EExecutionState *state, TranslationBlock *tb, uint64_t pc,
		unsigned vector) {
	if (vector != 0x80) { // in linux, 0x2e in windows
		return;
	}

	uint64_t pid = m_monitor->getPid(state);
	if (!m_process->isTracked(state, pid)) {
		return;
	}

	signal->connect(sigc::mem_fun(*this, &PythonS2E::onSyscallInt));
}

void PythonS2E::onSyscallInst(S2EExecutionState *state, uint64_t pc) {
	target_ulong syscallid = state->regs()->read<target_ulong>(
			CPU_OFFSET(regs[R_EAX]));
	processSyscallEntry(state, pc, syscallid);
}

void PythonS2E::onSyscallInt(S2EExecutionState *state, uint64_t pc) {
	target_ulong syscallid = state->regs()->read<target_ulong>(
			CPU_OFFSET(regs[R_EAX]));
	processSyscallEntry(state, pc, syscallid);
}

void PythonS2E::processSyscallEntry(S2EExecutionState *state, uint64_t pc,
		target_ulong syscallid) {
	uint64_t tid = m_monitor->getTid(state);
	auto it = m_sysCallRecord.find(tid);
	if (it != m_sysCallRecord.end()) {
		m_sysCallRecord.erase(it);
	}
	m_sysCallRecord.insert(std::make_pair(tid, (uint64_t) syscallid));

	s2e::plugins::callbacks::syscallEntry(state, (s2e::uint64) syscallid);
}
void PythonS2E::onModuleTranslateBlockStart(ExecutionSignal *signal,
		S2EExecutionState *state, const ModuleDescriptor &module,
		TranslationBlock *tb, uint64_t pc) {
	m_ins_connection_before.disconnect();
	m_ins_connection_after.disconnect();
	uint64_t pid = m_monitor->getPid(state);
	if (!m_process->isTracked(state, pid)) {
		return;
	}
	if (m_monitor->isKernelAddress(pc)) {
		return;
	}
	if (s2e::plugins::options::callbackAfter) {
		m_ins_connection_after =
				s2e()->getCorePlugin()->onTranslateInstructionEnd.connect(
						sigc::bind(
								sigc::mem_fun(*this,
										&PythonS2E::onTranslateInstructionEnd),
								(-module.LoadBase + module.NativeBase)));
	}

	if (s2e::plugins::options::callbackBefore) {
		m_ins_connection_before = s2e()->getCorePlugin()->onTranslateInstructionStart.connect(
				sigc::bind(
						sigc::mem_fun(*this,
								&PythonS2E::onTranslateInstructionStart),
						(-module.LoadBase + module.NativeBase)));
	}
}
void PythonS2E::onTranslateInstructionEnd(ExecutionSignal *signal,
		S2EExecutionState *state, TranslationBlock *tb, uint64_t pc,
		uint64_t addend) {
	//uint64_t modulePc = pc + addend;
	signal->connect(sigc::mem_fun(*this, &PythonS2E::afterExecuteInstruction));
}

void PythonS2E::onTranslateInstructionStart(ExecutionSignal *signal,
		S2EExecutionState *state, TranslationBlock *tb, uint64_t pc,
		uint64_t addend) {
	//uint64_t modulePc = pc + addend;
	signal->connect(sigc::mem_fun(*this, &PythonS2E::beforeExecuteInstruction));
}

void PythonS2E::beforeExecuteInstruction(S2EExecutionState *state,
		uint64_t pc) {
	unsigned char code[15];
	uint64_t codesize = 15;
	if (state->mem()->read(pc, code, codesize)) {
		s2e::plugins::Instruction *tritonInst = new s2e::plugins::Instruction();
		tritonInst->setOpcode(code, codesize);
		tritonInst->address = pc;
		s2e::plugins::callbacks::before(state, tritonInst);
	}
}
void PythonS2E::afterExecuteInstruction(S2EExecutionState *state, uint64_t pc) {
	unsigned char code[15];
	uint64_t codesize = 15;
	if (state->mem()->read(pc, code, codesize)) {
		s2e::plugins::Instruction *tritonInst = new s2e::plugins::Instruction();
		tritonInst->setOpcode(code, codesize);
		tritonInst->address = pc;
		s2e::plugins::callbacks::after(state, tritonInst);
	}
}

void PythonS2E::onEveryTranslateInstructionStart(ExecutionSignal *signal,
		S2EExecutionState *state, TranslationBlock *tb, uint64_t pc) {
	uint64_t pid = m_monitor->getPid(state);
	if (!m_process->isTracked(state, pid)) {
		return;
	}
	if (m_monitor->isKernelAddress(pc)) {
		return;
	}
	signal->connect(sigc::mem_fun(*this, &PythonS2E::onEveryExecInstruction),
			fsigc::signal_base::HIGHEST_PRIORITY);
}

void PythonS2E::onEveryExecInstruction(S2EExecutionState *state, uint64_t pc) {
	if (!m_sysCallRecord.size()) {
		return;
	}
	// check if we are after sysenter or syscall
	uint64_t tid = m_monitor->getTid(state);
	auto sysIt = m_sysCallRecord.find(tid);
	if (sysIt == m_sysCallRecord.end()) {
		return;
	}
	s2e::plugins::callbacks::syscallExit(state, (s2e::uint64) sysIt->second);
	m_sysCallRecord.erase(sysIt);
}

bool PythonS2E::execScript(const char *fileName) {
	/* On some Linux distro, we must load libpython to successfully load all others modules. See issue #276. */
	void *handle = dlopen(m_python_so_file.c_str(), RTLD_LAZY | RTLD_GLOBAL);
	if (!handle)
		throw std::runtime_error(
				"execScript(): Cannot load the Python library.");

	FILE *fd = nullptr;
	auto err = fopen_s(&fd, fileName, "r");
	if (err != 0)
		throw std::runtime_error("execScript(): Script file can't be found.");

	PyRun_SimpleFile(fd, fileName);

	fclose(fd);
	return true;
}

PythonS2E::~PythonS2E() {
	Py_Finalize();
	m_ins_connection_before.disconnect();
	m_ins_connection_after.disconnect();
}

int PythonS2E::fopen_s(FILE **fd, const char *fn, const char *flags) {
	*fd = fopen(fn, flags);
	if (*fd == 0)
		return -1;
	else
		return 0;
}
}
}
}
